import torch.nn as nn
import torch.nn.functional

class Net(nn.Module):
    def __init__(self,layers):
        super().__init__()
        self.fcnn=[]
        self.layers_num=len(layers)
        for i,j in zip(range(self.layers_num-1),range(1,self.layers_num)):
            self.fc=nn.Linear(layers[i], layers[j])
            self.fcnn.append(self.fc)
        self.fcnn = nn.ModuleList(self.fcnn)


    def forward(self, x):
        for i,fc in enumerate(self.fcnn):
            x = fc(x)
            if self.layers_num-1 !=i:
                x = nn.functional.relu(x)
        out = nn.functional.log_softmax(x, dim=1)  # softmax函数，dim=1表示按行计算概率分布
        return out



